import json
import pickle

import numpy as np
import scipy.io as sio


def _load(fp):
    """读取文件

    Args:
        fp: 路径

    Returns:
        数据
    """
    pos = fp.rfind('.')
    suffix = fp[pos + 1:]
    if suffix == 'npy':
        return np.load(fp)
    elif suffix == 'pkl':
        return pickle.load(open(fp, 'rb'))
    elif suffix == 'mat':
        return sio.loadmat(fp)


def save(dst, data, use_json=False):
    """保存文件为pkl和json

    Args:
        use_json: 是否生成json
        dst: 目标文件名称，不带后缀
        data: 要保存的np参数
    Returns:

    """
    if use_json:
        data_json = json.dumps(data.tolist())
        # json.dump(tri, StringIO('tri.json'))
        with open(dst + '.json', 'w') as f:
            f.write(data_json)
    with open(dst + '.pkl', 'wb') as f:
        pickle.dump(data, f)
